package org.example.brave.servlet;

import brave.Span;
import brave.SpanCustomizer;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.TraceContext;
import org.example.brave.http.HttpServerHandler;
import org.example.brave.http.HttpServerRequest;
import org.example.brave.http.HttpServerResponse;
import org.example.brave.http.HttpTracing;
import org.example.brave.servlet.internal.ServletRuntime;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.concurrent.atomic.AtomicBoolean;

public final class TracingFilter implements Filter {
    public static Filter create(Tracing tracing) {
        return new TracingFilter(HttpTracing.create(tracing));
    }

    public static Filter create(HttpTracing httpTracing) {
        return new TracingFilter(httpTracing);
    }

    final ServletRuntime servlet = ServletRuntime.get();
    final CurrentTraceContext currentTraceContext;
    final HttpServerHandler<HttpServerRequest, HttpServerResponse> handler;

    TracingFilter(HttpTracing httpTracing) {
        currentTraceContext = httpTracing.tracing().currentTraceContext();
        handler = HttpServerHandler.create(httpTracing);
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        HttpServletRequest req = (HttpServletRequest) request;
        HttpServletResponse res = servlet.httpServletResponse(response);

        // Prevent duplicate spans for the same request
        TraceContext context = (TraceContext) request.getAttribute(TraceContext.class.getName());
        if (context != null) {
            // A forwarded request might end up on another thread, so make sure it is scoped
            CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(context);
            try {
                chain.doFilter(request, response);
            } finally {
                scope.close();
            }
            return;
        }

        Span span = handler.handleReceive(new HttpServletRequestWrapper(req)); // span begin

        // Add attributes for explicit access to customization or span context
        request.setAttribute(SpanCustomizer.class.getName(), span.customizer());
        request.setAttribute(TraceContext.class.getName(), span.context());
        SendHandled sendHandled = new SendHandled();
        request.setAttribute(SendHandled.class.getName(), sendHandled);

        Throwable error = null;
        CurrentTraceContext.Scope scope = currentTraceContext.newScope(span.context());
        try {
            // any downstream code can see Tracer.currentSpan() or use Tracer.currentSpanCustomizer()
            chain.doFilter(req, res);
        } catch (Throwable e) {
            error = e;
            throw e;
        } finally {
            // When async, even if we caught an exception, we don't have the final response: defer
            if (servlet.isAsync(req)) {
                servlet.handleAsync(handler, req, res, span);
            } else if (sendHandled.compareAndSet(false, true)){
                // we have a synchronous response or error: finish the span
                HttpServerResponse responseWrapper = HttpServletResponseWrapper.create(req, res, error);
                handler.handleSend(responseWrapper, span); // span finish
            }
            scope.close();
        }
    }

    // Special type used to ensure handleSend is only called once
    static final class SendHandled extends AtomicBoolean {
    }

    @Override public void destroy() {
    }

    @Override public void init(FilterConfig filterConfig) {
    }
}
